# -*- coding: UTF-8 -*-
# cython: language_level=3
import multiprocessing
from concurrent import futures

from tqdm import tqdm

from base import base_parser
from base.base_parser import add_icp, fenlei
from db.mongodb import MongoDB
from utils.confident import site_checker
from utils.get_domain_info import get_website_info
from utils.ip_area import find_ip


def init_db():
    db = MongoDB()

    db.set_unique_key('todo_urls', 'host')
    db.set_ensure_index('todo_urls', 'status')

    db.set_unique_key('WEB_urls', 'host')
    db.set_ensure_index('WEB_urls', 'site_Confidence')
    db.set_ensure_index('WEB_urls', 'status')

    db.set_unique_key('WEB_search_info', 'url')
    db.set_ensure_index('WEB_search_info', 'Confidence')
    db.set_ensure_index('WEB_search_info', 'host')


def add_root_url(host, home_page, uid):
    website_info = get_website_info(home_page, host=host)
    ip = website_info['ip']

    try:
        assert ip
        country, province_name, city_name, isp = find_ip(ip)
        if country and country != '中国':
            province_name = country + province_name
    except:
        province_name = ''
        city_name = ''

    if host.endswith('.gov.cn') or host.endswith('.edu.cn') or host.endswith('.gov') or host.endswith('.edu'):
        status = 8
    else:
        if website_info['visit_status'] == 0:
            status = fenlei(website_info['title'], icp=website_info['icp'], province_name=province_name)
        else:
            status = 3

    add_icp(status=status, province_name=province_name, city_name=city_name, **website_info)
    MongoDB().update('todo_urls', {"_id": uid}, {"status": 2})


def add_root_url_t(INPUT_URLS, n=None):
    len_urls = len(INPUT_URLS)
    # print('添加url数量： ', len_urls)
    p_num = min(len_urls, 40)
    if not p_num:
        return
    db = MongoDB()

    tasks, results = [], []
    with futures.ThreadPoolExecutor(max_workers=p_num) as executor:
        for i in tqdm(INPUT_URLS, total=len_urls, desc='add_task'):
            db.update('todo_urls', {"_id": i['_id']}, {"status": 1})
            tasks.append(executor.submit(add_root_url, i['host'], i['home_page'], i['_id']))

        for task in tqdm(futures.as_completed(tasks), total=len_urls, desc=f'update_mongo-{n}'):
            # results.append(task.result())
            pass
    # print(results)


def add_root_url_main(max_workers=10):
    db = MongoDB()
    db.update('todo_urls', {"status": 1}, {"status": 0})
    INPUT_URLS = db.find('todo_urls', {'status': 0})
    len_urls = len(INPUT_URLS)
    print('添加url数量： ', len_urls, 'max_workers:', max_workers)

    if not len_urls:
        return
    step = int(len_urls / max_workers) + 1
    INPUT_URLS = [INPUT_URLS[i:i + step] for i in range(0, len_urls, step)]

    tasks, results = [], []
    with futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        for i, data in enumerate(INPUT_URLS):
            tasks.append(executor.submit(add_root_url_t, data, i))
        for task in tqdm(futures.as_completed(tasks), total=max_workers, desc='总进度'):
            # results.append(task.result())
            pass
    # print(results)


def check(item):
    base_parser.wuzheng_run('WEB_urls', item['host'])
    site_info = site_checker(item)
    try:
        site_confidence = site_info['site_Confidence']
        page_num = site_info['page_num']
    except:
        MongoDB().delete('WEB_search_info', {'host': item['host']})
    else:
        base_parser.wuzheng_update('WEB_urls', item['host'], 2, site_confidence, page_num)


def check_main(man_input=None):
    """
    检测主函数
    :param man_input: 手工导入mongo待检测数据，默认检测url表中status为0的数据
    """
    base_parser.clear_db()
    check_urls = base_parser.find_todo_url() if not man_input else man_input
    if not check_urls:
        return

    print('待检查url数量： ', len(check_urls))
    p_num = len(check_urls) if len(check_urls) < 16 else 16
    p = multiprocessing.Pool(p_num)
    print("-----start-----")
    for item in check_urls:
        p.apply_async(check, (item,))
    p.close()
    p.join()
    print("-----end-----")


if __name__ == "__main__":
    # init_db()

    # add_root_url('cnbz.gov.cn', 'https://www.cnbz.gov.cn', '1')
    # db.update('todo_urls', {}, {"status": 0})
    # db.delete('WEB_urls')

    add_root_url_main(max_workers=5)
    # check_main()

    # print(fenlei('深圳宣传片制作_首推万友影力传媒高端视频服务商_深圳企业宣传片制作_深圳产品宣传片拍摄_深圳影视广告_深圳微电影_深圳短视频拍摄'))
